{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd06b62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import BatchNorm2d\n",
    "\n",
    "affine_par = True\n",
    "\n",
    "\n",
    "def conv3x3(in_planes, out_planes, stride=1):\n",
    "    \"3x3 convolution with padding\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                     padding=1, bias=False)\n",
    "\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, fist_dilation=1, multi_grid=1):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n",
    "        self.bn1 = BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n",
    "                               padding=dilation * multi_grid, dilation=dilation * multi_grid, bias=False)\n",
    "        self.bn2 = BatchNorm2d(planes)\n",
    "        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n",
    "        self.bn3 = BatchNorm2d(planes * 4)\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "        self.relu_inplace = nn.ReLU(inplace=True)\n",
    "        self.downsample = downsample\n",
    "        self.dilation = dilation\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            residual = self.downsample(x)\n",
    "\n",
    "        out = out + residual\n",
    "        out = self.relu_inplace(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, layers):\n",
    "        self.inplanes = 128\n",
    "        super(ResNet, self).__init__()\n",
    "        self.conv1 = conv3x3(2, 64, stride=2)\n",
    "        self.bn1 = BatchNorm2d(64)\n",
    "        self.relu1 = nn.ReLU(inplace=False)\n",
    "        self.conv2 = conv3x3(64, 64)\n",
    "        self.bn2 = BatchNorm2d(64)\n",
    "        self.relu2 = nn.ReLU(inplace=False)\n",
    "        self.conv3 = conv3x3(64, 128)\n",
    "        self.bn3 = BatchNorm2d(128)\n",
    "        self.relu3 = nn.ReLU(inplace=False)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "\n",
    "        self.relu = nn.ReLU(inplace=False)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True)  # change\n",
    "        self.layer1 = self._make_layer(block, 64, layers[0])\n",
    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)\n",
    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4, multi_grid=(1, 1, 1))\n",
    "\n",
    "    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):\n",
    "        downsample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = nn.Sequential(\n",
    "                nn.Conv2d(self.inplanes, planes * block.expansion,\n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                BatchNorm2d(planes * block.expansion, affine=affine_par))\n",
    "\n",
    "        layers = []\n",
    "        generate_multi_grid = lambda index, grids: grids[index % len(grids)] if isinstance(grids, tuple) else 1\n",
    "        layers.append(block(self.inplanes, planes, stride, dilation=dilation, downsample=downsample,\n",
    "                            multi_grid=generate_multi_grid(0, multi_grid)))\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for i in range(1, blocks):\n",
    "            layers.append(\n",
    "                block(self.inplanes, planes, dilation=dilation, multi_grid=generate_multi_grid(i, multi_grid)))\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu1(self.bn1(self.conv1(x)))\n",
    "        x = self.relu2(self.bn2(self.conv2(x)))\n",
    "        x = self.relu3(self.bn3(self.conv3(x)))\n",
    "        x = self.maxpool(x)\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "        return x\n",
    "\n",
    "    def freeze(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.BatchNorm2d):\n",
    "                m.eval()\n",
    "\n",
    "\n",
    "class ResNetBackbone(nn.Module):\n",
    "\n",
    "    def __init__(self, pretrained=True):\n",
    "        super().__init__()\n",
    "        self.backbone = ResNet(Bottleneck, [3, 4, 23, 3])\n",
    "\n",
    "        if pretrained:\n",
    "            # saved_state_dict = torch.load('./network/pretrained_models/resnet101-imagenet.pth')\n",
    "            saved_state_dict = torch.load('/home/liuyang/Desktop/08_ice_tasks/03_公用模型/resnet101-imagenet.pth',\n",
    "                                          map_location=\"cpu\")\n",
    "            new_params = self.backbone.state_dict().copy()\n",
    "            for i in saved_state_dict:\n",
    "                if i!='conv1.weight':\n",
    "                    i_parts = i.split('.')\n",
    "                    if not i_parts[0] == 'fc':\n",
    "                        new_params['.'.join(i_parts[0:])] = saved_state_dict[i]\n",
    "\n",
    "            self.backbone.load_state_dict(new_params)\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.backbone(input)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
